package com.hunterai.comfyui;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.stream.Collectors;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;

@Component
public class ComfyUIWorkflowManager {
    
    private final ResourceLoader resourceLoader;
    private final ResourcePatternResolver resourcePatternResolver;
    private final ObjectMapper objectMapper;
    private final String workflowsBasePath;
    private final Map<String, JsonNode> workflowCache = new HashMap<>();
    private final Map<String, JsonNode> workflowParamsCache = new HashMap<>();
    
    public ComfyUIWorkflowManager(
            ResourceLoader resourceLoader, 
            ObjectMapper objectMapper,
            @Value("${comfyui.workflows.dir:classpath:workflows/}") String workflowsBasePath) {
        this.resourceLoader = resourceLoader;
        this.resourcePatternResolver = new PathMatchingResourcePatternResolver(resourceLoader);
        this.objectMapper = objectMapper;
        this.workflowsBasePath = workflowsBasePath;
        
        // 初始化时预加载工作流
        initializeWorkflows();
    }
    
    private void initializeWorkflows() {
        try {
            // 使用Spring的ResourcePatternResolver来查找所有匹配的资源
            Resource[] resources = resourcePatternResolver.getResources(workflowsBasePath + "*.json");
            
            for (Resource resource : resources) {
                // 排除参数配置文件（带有_params后缀的）
                String filename = resource.getFilename();
                if (filename != null && !filename.endsWith("_params.json")) {
                    String workflowId = filename.replace(".json", "");
                    
                    // 加载工作流
                    JsonNode workflow = objectMapper.readTree(resource.getInputStream());
                    workflowCache.put(workflowId, workflow);
                    
                    // 加载对应的参数配置
                    Resource paramsResource = resourceLoader.getResource(workflowsBasePath + workflowId + "_params.json");
                    if (paramsResource.exists()) {
                        JsonNode params = objectMapper.readTree(paramsResource.getInputStream());
                        workflowParamsCache.put(workflowId, params);
                    }
                }
            }
        } catch (IOException e) {
            throw new RuntimeException("Failed to initialize workflows", e);
        }
    }
    
    /**
     * 获取所有可用工作流ID
     */
    public List<String> getAvailableWorkflows() {
        return workflowCache.keySet().stream().collect(Collectors.toList());
    }
    
    /**
     * 获取指定工作流的参数配置
     */
    public JsonNode getWorkflowParams(String workflowId) {
        return workflowParamsCache.getOrDefault(workflowId, objectMapper.createObjectNode());
    }
    
    /**
     * 获取指定工作流
     */
    public JsonNode getWorkflow(String workflowId) {
        JsonNode workflow = workflowCache.get(workflowId);
        if (workflow == null) {
            throw new IllegalArgumentException("Workflow not found: " + workflowId);
        }
        return workflow.deepCopy(); // 返回副本以防止修改
    }
    
    /**
     * 应用用户参数到工作流
     */
    public JsonNode applyParams(String workflowId, Map<String, Object> userParams) {
        // 获取基础工作流
        JsonNode baseWorkflow = getWorkflow(workflowId);
        ObjectNode workflow = (ObjectNode) baseWorkflow;

            // 遍历所有节点找到KSampler节点并更新seed值
            workflow.fields().forEachRemaining(entry -> {
                // String nodeId = entry.getKey();
                JsonNode node = entry.getValue();
                
                // 检查是否是KSampler类型的节点
                if (node.has("class_type") && 
                    (node.get("class_type").asText().contains("KSampler") || 
                     node.get("class_type").asText().contains("SamplerCustom"))) {
                    
                    // 如果是KSampler节点，检查并修改inputs中的seed值
                    if (node.has("inputs") && node.get("inputs").has("seed")) {
                        // 生成随机种子
                        long randomSeed = new Random().nextInt(Integer.MAX_VALUE) + 1L;
                        
                        // 允许用户传入特定的seed值覆盖随机值
                        if (userParams.containsKey("seed")) {
                            randomSeed = Long.parseLong(userParams.get("seed").toString());
                        }
                        
                        // 更新节点的seed值
                        ((ObjectNode)node.get("inputs")).put("seed", randomSeed);
                    }
                }
            });        
        
        // 获取参数配置
        JsonNode paramsConfig = getWorkflowParams(workflowId);
        
        // 应用用户参数
        for (Map.Entry<String, Object> entry : userParams.entrySet()) {
            String paramKey = entry.getKey();
            Object paramValue = entry.getValue();
            
            // 查找参数配置
            if (paramsConfig.has(paramKey)) {
                JsonNode paramConfig = paramsConfig.get(paramKey);
                String nodeId = paramConfig.get("nodeId").asText();
                String paramPath = paramConfig.get("path").asText();
                
                // 获取节点
                if (workflow.has(nodeId)) {
                    ObjectNode node = (ObjectNode) workflow.get(nodeId);
                    
                    // 设置参数值
                    String[] pathParts = paramPath.split("\\.");
                    ObjectNode current = node;
                    
                    // 遍历路径直到最后一个部分
                    for (int i = 0; i < pathParts.length - 1; i++) {
                        if (current.has(pathParts[i])) {
                            current = (ObjectNode) current.get(pathParts[i]);
                        } else {
                            // 如果中间路径不存在，则创建
                            current = (ObjectNode) current.putObject(pathParts[i]);
                        }
                    }
                    
                    // 设置最终值
                    String lastPart = pathParts[pathParts.length - 1];
                    if (paramValue instanceof String) {
                        current.put(lastPart, (String) paramValue);
                    } else if (paramValue instanceof Integer) {
                        current.put(lastPart, (Integer) paramValue);
                    } else if (paramValue instanceof Long) {
                        current.put(lastPart, (Long) paramValue);
                    } else if (paramValue instanceof Double) {
                        current.put(lastPart, (Double) paramValue);
                    } else if (paramValue instanceof Boolean) {
                        current.put(lastPart, (Boolean) paramValue);
                    }
                }
            }
        }
        
        return workflow;
    }
}